// parser.c
#include "Yecc.h"

struct Token* cur;

struct Node* Program();

struct Node* StmtList();
struct Node* Stmt();

struct Node* ExprStmt();
struct Node* Expr();
struct Node* AsgExpr();         // includes AsgOp()
struct Node* CondExpr();
struct Node* LogExpr();
struct Node* BitExpr();
struct Node* CmpExpr();         // includes CmpOp()
struct Node* ShiftExpr();
struct Node* AddExpr();
struct Node* MulExpr();
struct Node* Primary();         // includes PrefixOp() and PostfixOp()
struct Node* Cast();
struct Node* PrimeExpr();       
struct Node* ArgList();

struct Node* CompoundStmt();
struct Node* IfStmt();
struct Node* SwitchStmt();
struct Node* CaseStmt();
struct Node* DefaultStmt();
struct Node* WhileStmt();
struct Node* ForStmt();
struct Node* DoStmt();
struct Node* BreakStmt();
struct Node* ContinueStmt();
struct Node* AsmStmt();
struct Node* ReturnStmt();

struct Node* DeclStmt();
struct Node* DeclType();
struct Type* PrimeType();
struct Type* ArraySuffix(struct Type* base);
struct Node* Initializer();
struct Node* InitList();

// parser tools
bool consume(const char* str);
void expect(const char* str);
bool peek(const char* str);

// consume() only consumes TK_PUNCT and TK_KEYWORD
bool consume(const char* str) {
    // Punct check
    if(cur->kind == TK_PUNCT) {
        enum PunctKind pk = getPunctKind(str);
        if(cur->pk == pk) {
            cur = cur->next;
            return true;
        }
    }
    // Keyword check
    if(cur->kind == TK_KEYWORD) {
        enum KeywordKind kw = getKeywordKind(str);
        if(cur->kw == kw && kw != KW_UNKNOWN) {
            cur = cur->next;
            return true;
        }
    }
    return false;
}

// it would consume a token
void expect(const char* str) {
    if(!consume(str)) {
        printf("Expected '%s' but ", str);
        printToken(cur);
        assert(0);
    }
}

// just peek, not consume
bool peek(const char* str) {
    bool PunctHit = (cur->kind == TK_PUNCT) && (cur->pk == getPunctKind(str));
    bool KeywordHit = (cur->kind == TK_KEYWORD) && (cur->kw == getKeywordKind(str));
    return PunctHit || KeywordHit;
}

// Peek next token (cur->next), not the current one.
// Return true if the next token matches the given string.
bool peek_next(const char* str) {
    if (!cur || !cur->next) return false;
    struct Token* nxt = cur->next;
    bool PunctHit = (nxt->kind == TK_PUNCT) && (nxt->pk == getPunctKind(str));
    bool KeywordHit = (nxt->kind == TK_KEYWORD) && (nxt->kw == getKeywordKind(str));
    return PunctHit || KeywordHit;
}

struct Node* new_node(enum NodeKind kind, struct Node* left, struct Node* right) {
    struct Node* node = (struct Node*)malloc(sizeof(struct Node));
    assert(node != NULL);
    node->kind = kind;
    node->left = left;
    node->right = right;
    node->condi = NULL;
    node->body = NULL;
    node->els = NULL;
    node->init = NULL;
    node->inc = NULL;
    node->type = NULL;
    node->num = 0;
    node->ch = '\0';
    node->str = NULL;
    node->name = NULL;
    node->next = NULL;
    return node;
}

struct Node* PrimeExpr() {
    // "(" Expr ")"
    if(consume("(")) {
        struct Node* node = Expr();
        expect(")");
        return node;
    }

    // "sizeof" "(" Type ")"
    if(consume("sizeof")) {
        expect("(");
        struct Node* type = DeclType();
        expect(")");
        struct Node* node = new_node(ND_SIZEOF, type, NULL);
        return node;
    }

    // number
    if(cur->kind == TK_NUM) {
        struct Node* node = new_node(ND_NUM, NULL, NULL);
        node->num = cur->num;
        cur = cur->next;
        return node;
    }

    // character
    if(cur->kind == TK_CH) {
        struct Node* node = new_node(ND_CH, NULL, NULL);
        node->ch = cur->ch;
        cur = cur->next;
        return node;
    }

    // string
    if(cur->kind == TK_STR) {
        struct Node* node = new_node(ND_STR, NULL, NULL);
        node->str = cur->str;
        cur = cur->next;
        return node;
    }

    // identifier
    if(cur->kind == TK_IDENT) {
        struct Node* node = new_node(ND_IDENT, NULL, NULL);
        node->name = cur->name;
        cur = cur->next;
        return node;
    }

    // true, false, NULL
    if(cur->kind == TK_KEYWORD) {
        struct Node* node = NULL;

        if(cur->kw == KW_TRUE) node = new_node(ND_TRUE, NULL, NULL);
        if(cur->kw == KW_FALSE) node = new_node(ND_FALSE, NULL, NULL);
        if(cur->kw == KW_NULL) node = new_node(ND_NULL, NULL, NULL);
        
        if(node == NULL) {
            printf("Unexpected keyword: ");
            printToken(cur);
            assert(0);
        }

        cur = cur->next;
        return node;
    }

    printf("Unexpected token: ");
    printToken(cur);
    assert(0);
}

struct Node* ArgList() {
    struct Node* node = AsgExpr();
    while(consume(",")) {
        node = new_node(ND_COMMA, node, AsgExpr());
    }
    return node;
}

bool isCast() {
    return peek("(") &&
         ( peek_next("void") 
        || peek_next("bool") 
        || peek_next("char") 
        || peek_next("short") 
        || peek_next("int") 
        || peek_next("long") 
        || peek_next("struct") 
        || peek_next("enum")
        || peek_next("signed")
        || peek_next("unsigned"));
}

struct Node* Cast() {
    expect("(");
    struct Node* node = DeclType();
    expect(")");
    return node;
}

struct Node* Primary() {
    if(consume("-")) {
        struct Node* zero = new_node(ND_NUM, NULL, NULL);
        struct Node* node = new_node(ND_SUB, zero, Primary());
        return node;
    }
    if(consume("!")) {
        return new_node(ND_LOG_NOT, Primary(), NULL);
    }
    if(consume("~")) {
        return new_node(ND_BITNOT, Primary(), NULL);
    }
    if(consume("&")) {
        return new_node(ND_ADDR, Primary(), NULL);
    }
    if(consume("*")) {
        return new_node(ND_DEREF, Primary(), NULL);
    }
    if(isCast()) {
        struct Node* type = Cast();
        return new_node(ND_CAST, type, Primary());
    }

    struct Node* node = PrimeExpr();

    while(true) {
        if(consume("++")) {
            node = new_node(ND_INC, node, NULL);
            continue;
        }
        if(consume("--")) {
            node = new_node(ND_DEC, node, NULL);
            continue;
        }
        // Function call
        if(consume("(")) {
            node = new_node(ND_FUNCALL, node, NULL);
            if (!consume(")")) {
                node->right = ArgList();
                expect(")");
            }
            continue;
        }
        // Array assess
        if (consume("[")) {
            struct Node* idx = Expr();
            expect("]");
            node = new_node(ND_INDEX, node, idx);
            continue;
        }
        // Struct member access by "." or "->"
        if (peek(".") || peek("->")){
            
            if(consume(".")) {
                node = new_node(ND_DOT, node, NULL);
            } else if(consume("->")) {
                node = new_node(ND_ARROW, node, NULL);
            } else {
                printf("Error: expected '.' or '->'\n");
                printToken(cur);
                assert(0);
            }
            
            if (cur->kind != TK_IDENT) {
                printf("Error: expected member name after '.' or '->'\n");
                printToken(cur);
                assert(0);
            }

            struct Node* ident = new_node(ND_IDENT, NULL, NULL);
            ident->name = cur->name;
            node->right = ident;
            cur = cur->next;
            continue;
        }

        return node;
    }
}

struct Node* MulExpr() {
    struct Node* node = Primary();
    while(true) {
        if (consume("*")) {
            node = new_node(ND_MUL, node, Primary());
        } else if (consume("/")) {
            node = new_node(ND_DIV, node, Primary());
        } else if (consume("%")) {
            node = new_node(ND_MOD, node, Primary());
        } else {
            return node;
        }
    }
}

struct Node* AddExpr() {
    struct Node* node = MulExpr();
    while(true) {
        if(consume("+")) {
            node = new_node(ND_ADD, node, MulExpr());
        } else if(consume("-")) {
            node = new_node(ND_SUB, node, MulExpr());
        } else {
            return node;
        }
    }
}

struct Node* ShiftExpr() {
    struct Node* node = AddExpr();
    while(true) {
        if(consume("<<")) {
            node = new_node(ND_SHL, node, AddExpr());
        } else if(consume(">>")) {
            node = new_node(ND_SHR, node, AddExpr());
        } else {
            return node;
        }
    }
}

struct Node* CmpExpr() {
    struct Node* node = ShiftExpr();
    while(true) {
        if(consume("==")) {
            node = new_node(ND_EQ, node, ShiftExpr());
        } else if(consume("!=")) {
            node = new_node(ND_NEQ, node, ShiftExpr());
        } else if(consume("<=")) {
            node = new_node(ND_LE, node, ShiftExpr());
        } else if(consume(">=")) {
            node = new_node(ND_GE, node, ShiftExpr());
        } else if(consume("<")) {
            node = new_node(ND_LT, node, ShiftExpr());
        } else if(consume(">")) {
            node = new_node(ND_GT, node, ShiftExpr());
        } else {
            return node;
        }
    }
}

struct Node* BitExpr() {
    struct Node* node = CmpExpr();
    while(true) {
        if(consume("&")) {
            node = new_node(ND_BITAND, node, CmpExpr());
        } else if(consume("|")) {
            node = new_node(ND_BITOR, node, CmpExpr());
        } else if(consume("^")) {
            node = new_node(ND_BITXOR, node, CmpExpr());
        } else {
            return node;
        }
    }
}

struct Node* LogExpr() {
    struct Node* node = BitExpr();
    while(true) {
        if(consume("&&")) {
            node = new_node(ND_LOG_AND, node, BitExpr());
        } else if(consume("||")) {
            node = new_node(ND_LOG_OR, node, BitExpr());
        } else {
            return node;
        }
    }
}

struct Node* CondExpr() {
    struct Node* node = LogExpr();
    if(consume("?")) {
        struct Node* then_expr = Expr();
        expect(":");
        struct Node* else_expr = CondExpr();

        struct Node* tri_expr = new_node(ND_COND, NULL, NULL);
        tri_expr->condi = node;
        tri_expr->left  = then_expr;
        tri_expr->right = else_expr;
        return tri_expr;
    }
    return node;
}

struct Node* AsgExpr() {
    struct Node* node = CondExpr();
    // Assignment operators
    if(consume("="))   return new_node(ND_ASSIGN, node, AsgExpr());
    if(consume("+="))  return new_node(ND_ADD_ASG, node, AsgExpr());
    if(consume("-="))  return new_node(ND_SUB_ASG, node, AsgExpr());
    if(consume("*="))  return new_node(ND_MUL_ASG, node, AsgExpr());
    if(consume("/="))  return new_node(ND_DIV_ASG, node, AsgExpr());
    if(consume("%="))  return new_node(ND_MOD_ASG, node, AsgExpr());
    if(consume("&="))  return new_node(ND_AND_ASG, node, AsgExpr());
    if(consume("|="))  return new_node(ND_OR_ASG, node, AsgExpr());
    if(consume("^="))  return new_node(ND_XOR_ASG, node, AsgExpr());
    if(consume("<<=")) return new_node(ND_SHL_ASG, node, AsgExpr());
    if(consume(">>=")) return new_node(ND_SHR_ASG, node, AsgExpr());
    return node;
}

struct Node* Expr() {
    struct Node* node = AsgExpr();
    if(consume(","))
        node = new_node(ND_COMMA, node, Expr());
    return node;
}

struct Node* ExprStmt() {
    struct Node* node = NULL;
    if(!peek(";")) node = Expr();
    expect(";");
    return node;
}

struct Node* CompoundStmt() {
    expect("{");
    struct Node* node = StmtList();
    expect("}");
    return node;
}
struct Node* IfStmt() {
    expect("if");
    expect("(");
    struct Node* condi = Expr();
    expect(")");

    struct Node* body = Stmt();
    struct Node* els = NULL;

    if(consume("else")) els = Stmt();

    struct Node* node = new_node(ND_IF, NULL, NULL);
    node->condi = condi;
    node->body = body;
    node->els = els;
    return node;
}

struct Node* WhileStmt() {
    expect("while");
    expect("(");
    struct Node* condi = Expr();
    expect(")");

    struct Node* body = Stmt();

    struct Node* node = new_node(ND_WHILE, NULL, NULL);
    node->condi = condi;
    node->body = body;
    return node;
}

struct Node* DoStmt() {
    expect("do");
    struct Node* body = Stmt();
    expect("while");
    expect("(");
    struct Node* condi = Expr();
    expect(")");
    expect(";");

    struct Node* node = new_node(ND_DO, NULL, NULL);
    node->body = body;
    node->condi = condi;
    return node;
}

struct Node* ForStmt() {

    struct Node* init = NULL;
    struct Node* condi = NULL;
    struct Node* inc = NULL;

    expect("for");
    expect("(");
    if(!peek(";")) init = Expr();
    expect(";");
    if(!peek(";")) condi = Expr();
    expect(";");
    if(!peek(")")) inc = Expr();
    expect(")");

    struct Node* body = Stmt();

    struct Node* node = new_node(ND_FOR, NULL, NULL);
    node->init = init;
    node->condi = condi;
    node->inc = inc;
    node->body = body;
    return node;
}

struct Node* BreakStmt() {
    expect("break");
    struct Node* node = new_node(ND_BREAK, NULL, NULL);
    expect(";");
    return node;
}

struct Node* ContinueStmt() {
    expect("continue");
    struct Node* node = new_node(ND_CONTINUE, NULL, NULL);
    expect(";");
    return node;
}

struct Node* CaseStmt() {
    expect("case");
    struct Node* condi = Expr();
    expect(":");
    struct Node* body = Stmt();

    struct Node* node = new_node(ND_CASE, NULL, NULL);
    node->condi = condi;
    node->body = body;
    return node;
}

struct Node* DefaultStmt() {
    expect("default");
    expect(":");
    struct Node* body = Stmt();

    struct Node* node = new_node(ND_DEFAULT, NULL, NULL);
    node->body = body;
    return node;
}

struct Node* AsmStmt() {
    expect("asm");
    expect("(");
    assert(cur->kind == TK_STR);
    struct Node* body = new_node(ND_STR, NULL, NULL);
    body->str = cur->str;
    cur = cur->next;
    expect(")");
    expect(";");

    struct Node* node = new_node(ND_ASM, NULL, NULL);
    node->body = body;

    return node;
}

struct Node* ReturnStmt() {
    expect("return");
    struct Node* body = ExprStmt();

    struct Node* node = new_node(ND_RETURN, NULL, NULL);
    node->body = body;
    return node;
}

struct Node* SwitchStmt() {
    expect("switch");
    expect("(");
    struct Node* condi = Expr();
    expect(")");
    struct Node* body = Stmt();

    struct Node* node = new_node(ND_SWITCH, NULL, NULL);
    node->condi = condi;
    node->body = body;
    return node;
}

bool isType() {
    return peek("void") 
        || peek("bool") 
        || peek("char") 
        || peek("short") 
        || peek("int") 
        || peek("long") 
        || peek("struct") 
        || peek("enum");
}

struct Type* PrimeType() {

    struct Type* type = (struct Type* )malloc(sizeof(struct Type));
    assert(type != NULL);
    type->kind = TY_UNKNOWN;
    type->base = NULL;
    type->len = 0;
    type->name = NULL;
    type->sign = SG_UNKNOWN;

    bool tag = false;
    if(consume("void")) type->kind = TY_VOID;
    if(consume("bool")) type->kind = TY_BOOL;
    if(consume("char")) type->kind = TY_CHAR;
    if(consume("short")) type->kind = TY_SHORT;
    if(consume("int")) type->kind = TY_INT;
    if(consume("long")) type->kind = TY_LONG;
    if(consume("struct")) {
        type->kind = TY_STRUCT;
        tag = true;
    }
    if(consume("enum")) {
        type->kind = TY_ENUM;
        tag = true;
    }

    assert(type->kind != TY_UNKNOWN);

    if(tag) {
        type->name = cur->name;
        cur = cur->next;
    }

    return type;
}

// DeclType = ("signed" | "unsigned")? PrimeType ("*")*
struct Node* DeclType() {

    enum SignKind sign = SG_UNKNOWN;
    if(consume("signed")) sign = SG_SIGNED;
    else if(consume("unsigned")) sign = SG_UNSIGNED;

    struct Type* type = PrimeType();
    type->sign = sign;

    while(consume("*")) {
        struct Type* ptrType = (struct Type*)malloc(sizeof(struct Type));
        assert(ptrType != NULL);
        ptrType->kind = TY_PTR;
        ptrType->sign = SG_UNKNOWN;
        ptrType->len = 0;
        ptrType->name = NULL;
        ptrType->base = type;   // point to the previous one
        type = ptrType;         // update the current one
    }

    struct Node* node = new_node(ND_TYPE, NULL, NULL);
    node->type = type;
    return node;
}

struct Type* ArraySuffix(struct Type* base) {

    struct Type* type = base;

    while(consume("[")) {

        long len = 0;
        if(cur->kind == TK_NUM) {
            len = cur->num;
            cur = cur->next;
        }
        expect("]");

        struct Type* arrType = (struct Type*)malloc(sizeof(struct Type));
        assert(arrType != NULL);
        arrType->kind = TY_ARY;
        arrType->sign = SG_UNKNOWN;
        arrType->len = len;
        arrType->name = NULL;
        arrType->base = type;

        type = arrType;
    }

    return type;
}

struct Node* Initializer() {
    if(peek("{")) return InitList();
    else return AsgExpr();
}

struct Node* InitList() {
    expect("{");

    struct Node* head = new_node(ND_INILIST, NULL, NULL);
    head->init = Initializer();
    struct Node* curNode = head;

    while(consume(",")) {
        struct Node* nextItem = new_node(ND_INILIST, NULL, NULL);
        nextItem->init = Initializer();
        curNode->next = nextItem;
        curNode = nextItem;
    }
    expect("}");
    return head;
}

// DeclStmt = DeclType Ident (ArraySuffix)* ("=" Initializer)? ";"
struct Node* DeclStmt() {

    struct Node* typeNode = DeclType();
    assert(typeNode != NULL);

    assert(cur->kind == TK_IDENT);
    struct Node* ident = PrimeExpr();
    assert(ident->kind == ND_IDENT);

    typeNode->type = ArraySuffix(typeNode->type);

    struct Node* init = NULL;
    if(consume("=")) init = Initializer();
    
    expect(";");

    struct Node* node = new_node(ND_DECL, typeNode, ident);
    node->init = init;

    return node;
}

struct Node* Stmt() {

    if(peek("if")) return IfStmt();
    if(peek("switch")) return SwitchStmt();
    if(peek("case")) return CaseStmt();
    if(peek("default")) return DefaultStmt();
    
    if(peek("while")) return WhileStmt();
    if(peek("do")) return DoStmt();
    if(peek("for")) return ForStmt();
    if(peek("break")) return BreakStmt();
    if(peek("continue")) return ContinueStmt();

    if(peek("{")) return CompoundStmt();
    if(peek("asm")) return AsmStmt();
    if(peek("return")) return ReturnStmt();

    if(isType()) return DeclStmt();
    
    return ExprStmt();
}

// Parse multiple statements into a linked list via node->next
struct Node* StmtList() {
    struct Node* head = NULL;
    struct Node* tail = NULL;
    while(true) {

        if(cur->kind == TK_EOF) break;
        if(peek("}")) break;

        struct Node* stmt = Stmt();

        if(!head) head = stmt;
        else tail->next = stmt;

        tail = stmt;
    }
    return head;
}

struct Node* Program() {
    return StmtList();
}

void FreeToken(struct Token* head) {
    struct Token* curTok = head;
    while(curTok != NULL) {
        struct Token* nextTok = curTok->next;
        free(curTok->name);
        free(curTok->str);
        free(curTok);
        curTok = nextTok;
    }
    return ;
}

struct Node* Parser(struct Token* head) {

    cur = head->next;

    struct Node* root = Program();

    printf("\n### Parser Completed! ###\n\n");

    debug_info_dumpAbstractSyntaxTree(root);

    FreeToken(head);

    return root;
}
